package learning.data;

import java.util.Arrays;
import java.util.Comparator;
import java.util.HashMap;
import java.util.Iterator;
import java.util.Random;

import learning.data.document.IDocument;

public class Dataset<T extends IDocument> implements Iterable<T> {

	private T[] docs;
	private Random random;
	private String[] featureNames;
	private HashMap<String,Integer> featureIds;
	private String[] labelNames;
	private HashMap<String,Integer> labelIds;
	private int numLabels;
	private int numFeatures;
	
	private int viewStart = -1;
	private int viewEnd = -1;
	private boolean viewInvert = false;
	
	Dataset() {
		random = new Random(System.currentTimeMillis());
	}
	
	Dataset(T[] docs, String[] featureNames, HashMap<String, Integer> featureIds,			
			String[] labelNames, HashMap<String, Integer> labelIds) {
		this();
		this.docs = docs;
		this.featureNames = featureNames;
		this.featureIds = featureIds;
		this.labelNames = labelNames;
		this.labelIds = labelIds;
		this.numLabels = labelNames.length;
		//this.numFeatures = featureNames.length;
		this.numFeatures = featureIds.size();
	}

	// getters
	// returns document at index, ignoring views
	public T doc(int index) { return docs[index]; }
	
	public int numDocs() { return docs.length; }
	
	public int numLabels() { return numLabels; }

	public int numFeatures() { return numFeatures; }

	public String[] getFeatureNames() { return featureNames; }
	
	public HashMap<String, Integer> getFeatureIds() { return featureIds; }
	
	public String[] getLabelNames() { return labelNames; }
	
	public HashMap<String, Integer> getLabelIds() { return labelIds; } 
	
	public String getFeatureName(int id) { return featureNames[id]; }
	
	public int getFeatureId(String name) { return featureIds.get(name); }
	
	public String getLabelName(int id) { return labelNames[id]; }
	
	public int getLabelId(String name) { return labelIds.get(name); }

	// shuffle
	public void randomizeOrder() {
		
		if (viewStart < 0)
			for (int i=0; i < docs.length; i++) {
				// pick element that we want to swap with
				int e = i + random.nextInt(docs.length - i);
				T tmp = docs[e];
				docs[e] = docs[i];
				docs[i] = tmp;
			}
		else {
			// only randomize order on view, if view is defined
			if (!viewInvert) {						
				for (int i=viewStart; i < viewEnd; i++) {
					int e = i + random.nextInt(viewEnd - viewStart - i);
					T tmp = docs[e];
					docs[e] = docs[i];
					docs[i] = tmp;
				}
			} else {
				int total = docs.length - (viewEnd - viewStart);
				for (int i=0; i < total; i++) {
					int e = i + random.nextInt(total - i);
					
					int it = i, et = e;
					if (i >= viewStart) it = i + (viewEnd - viewStart);
					if (e >= viewStart) et = e + (viewEnd - viewStart);
					
					T tmp = docs[et];
					docs[et] = docs[it];
					docs[it] = tmp;
				}
			}
		}
	}

	// views
	public void releaseIteratorView() {
		this.viewStart = -1;
		this.viewEnd = -1;
	}
	
	public void setIteratorView(int numBlocks, int selectedBlock, boolean invert) {
		viewStart = (numDocs() / numBlocks) * selectedBlock + Math.min(selectedBlock, numDocs() % numBlocks);
		viewEnd = (numDocs() / numBlocks) * (selectedBlock+1) + Math.min(selectedBlock+1, numDocs() % numBlocks);
		viewInvert = invert;
	}
	
	public int numDocsInView() {
		if (!viewInvert) { return viewEnd - viewStart; }
		else { return docs.length - viewEnd + viewStart; }
	}
	
	public Iterator<T> iterator() {
		if (this.viewStart < 0)		
			// case 1: no view selection
			return new Iterator<T>() {
				int cursor = 0;
				
				public boolean hasNext() {
					return (cursor < docs.length);
				}
	
				public T next() {
					return docs[cursor++];
				}
	
				public void remove() { }
			};
		else
			// case 2: view selection
			return new Iterator<T>() {
				int cursor = viewInvert? 
						((0 < viewStart)? 0 : viewEnd) : viewStart;
				
				public boolean hasNext() {
					if (!viewInvert)
						return viewStart <= cursor && cursor < viewEnd;
					else
						return (0 <= cursor && cursor < viewStart) ||
							   (viewEnd <= cursor && cursor < docs.length);
				}

				public T next() {
					if (!viewInvert)
						return docs[cursor++];
					else {
						T d = docs[cursor++];
						if (cursor == viewStart) cursor = viewEnd;
						return d;
					}
				}

				public void remove() { }			
			};
	}	
	
	// sometimes, some of the examples in a dataset are dependent on each other (e.g. drawn from the same larger document)
	// if these have the same dependencyAttribute, they will always be grouped into one block
	public void setIteratorView(int numBlocks, int selectedBlock, boolean invert, final int dependencyAttribute) {
		Arrays.sort(docs, new Comparator<T>() {
			public int compare(T doc1, T doc2) {
				return doc1.getDependencyAttribute(dependencyAttribute) -
				       doc2.getDependencyAttribute(dependencyAttribute);
			}});
		
		// computing the start and end of each block is a bit more challenging now
		// ... TODO
		
		viewStart = (numDocs() / numBlocks) * selectedBlock + Math.min(selectedBlock, numDocs() % numBlocks);
		viewEnd = (numDocs() / numBlocks) * (selectedBlock+1) + Math.min(selectedBlock+1, numDocs() % numBlocks);
		viewInvert = invert;
	}

}
